[VectorDistribute] Allow duplication of operation chains#23937
[VectorDistribute] Allow duplication of operation chains#23937sommerlukas wants to merge 3 commits intoiree-org:mainfrom
Conversation
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
sommerlukas
requested review from
MaheshRavishankar,
Max191,
hanhanW and
qedawkins
as code owners
compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis_chain_cloning.mlir
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
Outdated
Show resolved
Hide resolved
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
Groverkss
left a comment
Groverkss
left a comment
There was a problem hiding this comment.
LGTM, feel free to land after addressing comments
|
|
||
| /// Maximum length of chains of cheap-to-compute operations that get duplicated | ||
| /// for layout conflict resolution. | ||
| static constexpr unsigned kMaxChainLength = 8; |
There was a problem hiding this comment.
any reason for unsigned? Can we just use int? (No need to reply, just keep it if you have a reason)
There was a problem hiding this comment.
The only place we use it is for a comparison with size(), so I now decided to make it size_t to match the return type.
| static bool isCheapToClone(Operation *op) { | ||
| if (isDuplicatableLeaf(op)) { | ||
| return true; | ||
| } | ||
| return isPure(op) && | ||
| (isa<vector::BroadcastOp, vector::TransposeOp, vector::ShapeCastOp>( | ||
| op) || | ||
| OpTrait::hasElementwiseMappableTraits(op)); | ||
| } | ||
|
|
||
| /// Collect a chain of ops that can be cloned together. Starting from `op`, | ||
| /// walk backward through single-result, cheap-to-clone ops until we reach | ||
| /// duplicatable leaves, constants, or non-vector operands. Returns true if | ||
| /// the entire chain is safe to clone. Shared intermediates (with multiple | ||
| /// uses) are allowed because all ops in the chain are cheap to duplicate. |
There was a problem hiding this comment.
Not sure why single result is important here. We happen to have most things as single result today, but just curious why.
There was a problem hiding this comment.
The comment was outdated, the code didn't check for a single result anymore, but I forgot to update the comment.
| // Non-vector operands (scalars, indices) don't need cloning. | ||
| if (!isa<VectorType>(operand.getType())) { | ||
| continue; | ||
| } |
There was a problem hiding this comment.
We could probably say this for any vector with 1 element as well.
There was a problem hiding this comment.
Done, I've added an early continue for single-element vectors.
| for (Operation *op : llvm::reverse(chain)) { | ||
| fixupOp(mapping.lookup(op->getResult(0)).getDefiningOp()); | ||
| } |
There was a problem hiding this comment.
Can we add a comment on why this will not recursively loop?
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
3 participants
If a value gets assigned two different layouts, the VectorLayoutAnalysis tries to duplicate that value if it is easily duplicatable to avoid layout transformation via LDS.
So far, this duplication was limited to single operations. This PR extends this to bounded chains of easy to duplicate/recompute values.
The motivation for this change is masks. CSE will deduplicate
create_maskoperations. After early materialization ofcreate_maskoperations, the mask is transformed to a chain of operations:As the mask is used for different operations, it will receive multiple different layouts. Therefore, we would materialize these masks to LDS for layout transformation.
By allowing chains of cheap operations to be duplicated, we avoid that materialization in LDS.
This is part of #23415.
Assisted-by: Claude Code